from getdata import importData
import os, math
import numpy as np

PATH = os.path.dirname(os.path.abspath(__file__))  # 当前文件夹路径
DATASET_PATH = PATH + '/folder'  # 数据文件位置
continuous_col = [0, 3, 4, 7, 9] # 数据为连续型的列


def calc_prior_prob(label):
    # 求先验概率
    prior_prob = {}
    Y = set(label)
    for y in Y:
        prior_prob[y] = len([i for i in label if i == y]) / len(label)
    return prior_prob


def calc_likelihood(data, label, col):
    """
    求似然估计，离散属性、连续属性
    :param data:
    :param label:
    :return:
    """
    likelihood_prob = {}
    X = list(set(data))
    Y = list(set(label))

    if col in continuous_col:  # 数据连续分布
        for i in range(len(Y)):
            divide = [(x, y) for x, y in zip(data, label) if y == Y[i]]
            # print(divide)
            mean = np.mean(divide[0])  # 数据列的均值
            std = np.std(divide[0])  # 数据列的标准差
            for x_i in X:
                # 高斯分布计算连续分布的属性概率
                likelihood_prob[(x_i, Y[i])] = 1 / (math.sqrt(2 * math.pi) * std) * math.exp((-(x_i - mean) ** 2) / (2 * std * std))
    else:  # 数据离散分布
        for i in range(len(Y)):
            divide = [(x, y) for x, y in zip(data, label) if y == Y[i]]
            # print(divide)
            for x_i in X:
                count = 0
                for num in range(len(divide)):
                    if x_i == divide[num][0]:
                        count += 1
                likelihood_prob[(x_i, Y[i])] = count / len(divide)  # 未加入拉普拉斯平滑
    return likelihood_prob


def classfier(data, label):
    # 分类器
    _, col = np.shape(data)
    label = [int(i) for i in label[:, 0]]
    prior_prob = calc_prior_prob(label)  # 先验概率字典
    ans = []
    for i in range(col):
        # print(data[:, i:i + 1][:, 0])
        ans.append(calc_likelihood(data[:, i:i + 1][:, 0], label, i))  # 求每一个属性的似然估计

    return prior_prob, ans


def predict_one(data, prior_prob, likelihood):
    ans1 = math.log2(prior_prob[1])
    ans2 = math.log2(prior_prob[2])
    for i in range(len(data)):
        if (data[i], 1) in likelihood[i].keys():
            ans1 += math.log2(likelihood[i][(data[i], 1)])
        if (data[i], 2) in likelihood[i].keys():
            ans2 += math.log2(likelihood[i][(data[i], 2)])
    return 1 if ans1 > ans2 else 2


def predict(data, prior_prob, likelihood):
    predict_labels = []
    for item in data:
        predict_labels.append(predict_one(list(item), prior_prob, likelihood))
    return predict_labels


def accuracy(real_label, predict_label):
    count = 0
    for i in range(len(real_label)):
        if predict_label[i] == real_label[i]:
            count += 1
    print(count, len(real_label))
    return count / len(real_label)


def main(filename):
    data = np.array(importData(DATASET_PATH, filename))  # 获取数据集

    train_data = data[:, 1:]
    train_label = data[:, :1]

    # 分类器
    prior_prob, likelihood_prob = classfier(train_data, train_label)
    print(likelihood_prob)
    # 预测
    predict1 = predict_one([61.0,1.0,1.0,134.0,234.0,0.0,0.0,145.0,0.0,2.6,2.0,2.0,3.0], prior_prob, likelihood_prob)
    print(predict1)

    predict_labels = predict(train_data, prior_prob, likelihood_prob)

    # 准确率
    print(predict_labels)
    print([int(i) for i in train_label[:, 0]])
    print(">>> 朴素贝叶斯的准确率为%f" % accuracy([int(i) for i in train_label[:, 0]], predict_labels))


if __name__ == '__main__':
    main(DATASET_PATH + '/heart.data')